from ppci import ir
from ppci import irutils
import logging
import io

logging.basicConfig(level=logging.WARNING, format="%(levelname)s - %(message)s")
logger = logging.getLogger("interpreter")
logger.setLevel(logging.DEBUG)


class ExecutionContext:
    """执行上下文，保存当前函数、基本块、指令指针和变量值"""

    def __init__(self, function, args):
        self.function = function
        self.variables = {}
        self.block = function.entry
        self.ip = 0
        self.return_value = None
        self.previous_block = None
        self.caller_return_var = None  # 存储调用者指定的返回变量名

        # 初始化参数
        for param, arg in zip(function.arguments, args):
            self.variables[param.name] = arg


class Interpreter:
    """PPCI IR解释器"""

    def __init__(self):
        self.call_stack = []
        self.memory = {}  # 模拟内存，键为地址（整数），值为存储的值
        self.globals = {}
        self.current_ctx = None
        self.next_address = 0x1000  # 初始内存地址
        self.allocations = {}  # 记录Alloc变量名对应的基地址

    def push_context(self, function, args, return_var=None):
        """压入新的执行上下文"""
        ctx = ExecutionContext(function, args)
        ctx.caller_return_var = return_var  # 记录调用者需要的返回变量名
        self.call_stack.append(ctx)
        self.current_ctx = ctx
        self.current_ctx.ip = 0

    def pop_context(self):
        """弹出当前上下文并传递返回值"""
        if self.call_stack:
            popped = self.call_stack.pop()
            # 将返回值传递给上层调用者
            if self.call_stack:
                caller_ctx = self.current_ctx = self.call_stack[-1]
                if popped.return_value is not None and popped.caller_return_var:
                    caller_ctx.variables[popped.caller_return_var] = popped.return_value
            return popped.return_value

    def get_value(self, value):
        """获取值（常量或变量）"""
        if isinstance(value, ir.Const):
            return value.value
        elif isinstance(value, ir.Parameter):
            return self.current_ctx.variables.get(value.name, 0)
        return self.current_ctx.variables.get(getattr(value, "name", ""), 0)

    def run(self, module):
        """运行模块"""
        self.globals = {f.name: f for f in module.functions}
        main_func = module.get_function("main")
        self.push_context(main_func, [])

        while self.current_ctx:
            return_val = self.execute_instruction()
            if return_val is not None:
                return return_val

    def execute_instruction(self):
        """指令执行分发器"""

        instr = self.current_ctx.block.instructions[self.current_ctx.ip]
        self.current_ctx.ip += 1
        # 指令类型映射
        handlers = {
            ir.Const: self.execute_const,
            ir.LiteralData: self.execute_literal_data,
            ir.Binop: self.execute_binop,
            ir.FunctionCall: self.execute_function_call,
            ir.ProcedureCall: self.execute_procedure_call,
            ir.Return: self.execute_return,
            ir.CJump: self.execute_cjump,
            ir.Jump: self.execute_jump,
            ir.Phi: self.execute_phi,
            ir.Alloc: self.execute_alloc,
            ir.AddressOf: self.execute_address_of,
            ir.Load: self.execute_load,
            ir.Store: self.execute_store,
            ir.Cast: self.execute_cast,
            ir.Exit: self.execute_exit,
        }

        handler = handlers.get(type(instr))
        if handler:
            ret_val = handler(instr)
            return ret_val
        else:
            raise RuntimeError(f"Unsupported instruction: {type(instr).__name__}")

    def execute_const(self, instr):
        """处理常量指令"""
        self.current_ctx.variables[instr.name] = instr.value

    def execute_literal_data(self, instr):
        """处理字面量数据指令"""
        # 分配内存地址
        base_address = self.next_address
        # 计算对齐后的地址（这里假设按1字节对齐）
        self.next_address += instr.ty.size
        # 将字节数据写入内存
        for offset, byte in enumerate(instr.data):
            self.memory[base_address + offset] = byte
        # 记录变量名对应的基地址
        self.allocations[instr.name] = base_address
        self.current_ctx.variables[instr.name] = base_address

    def execute_binop(self, instr):
        a = self.get_value(instr.a)
        b = self.get_value(instr.b)
        op_map = {
            "+": lambda x, y: x + y,
            "-": lambda x, y: x - y,
            "*": lambda x, y: x * y,
            "/": lambda x, y: x // y,
        }
        self.current_ctx.variables[instr.name] = op_map[instr.operation](a, b)

    def _read_cstring(self, address):
        """从内存读取以null结尾的字符串"""
        chars = []
        while True:
            char = self.memory.get(address, 0)
            if char == 0:
                break
            chars.append(chr(char))
            address += 1
        return "".join(chars)

    def execute_external_funcation(self, callee_name, args):
        def printf_impl(args):
            # 读取格式字符串
            format_ptr = args[0]
            fmt_str = self._read_cstring(format_ptr)

            # 读取可变参数
            va_args = []
            for va_ptr in args[1:]:
                # 假设按i32类型读取参数
                value = self.memory.get(va_ptr, 0)
                va_args.append(value)

            # 模拟printf输出
            try:
                print(fmt_str % tuple(va_args), end="")
            except:
                print(f"<格式化错误: {fmt_str} {va_args}>", end="")
            # 假设返回打印字符数（这里简化返回0）
            return 0

        function_impls = {
            "printf": printf_impl,
        }
        function_impl = function_impls[callee_name]
        return function_impl(args)

    def execute_function_call(self, instr):
        """处理函数调用"""
        callee = self.globals.get(instr.callee.name)
        args = [self.get_value(arg) for arg in instr.arguments]
        # 处理外部函数调用
        if not callee:
            self.current_ctx.variables[instr.name] = self.execute_external_funcation(
                instr.callee.name, args
            )
            return
        # 保存当前上下文的位置
        saved_block = self.current_ctx.block
        saved_ip = self.current_ctx.ip - 1  # 回退到调用指令的位置

        # 压入新上下文
        self.push_context(callee, args, return_var=instr.name)

        # 执行被调用函数
        while self.current_ctx and self.current_ctx.function == callee:
            return_val = self.execute_instruction()

        # 恢复调用者的执行位置
        self.current_ctx.block = saved_block
        self.current_ctx.ip = saved_ip + 1  # 执行下一条指令

        # 显式存储返回值到变量
        if return_val is not None:
            self.current_ctx.variables[instr.name] = return_val

    def execute_procedure_call(self, instr):
        callee = self.globals.get(instr.callee.name)
        args = [self.get_value(arg) for arg in instr.arguments]
        self.push_context(callee, args)
        while self.current_ctx and self.current_ctx.function == callee:
            self.execute_instruction()

    def execute_return(self, instr):
        """处理返回指令"""
        self.current_ctx.return_value = self.get_value(instr.result)
        if self.current_ctx.function.name == 'main':
            self.current_ctx = None
        return self.pop_context()  # 这里直接返回并弹出上下文

    def execute_cjump(self, instr):
        """处理条件跳转"""
        a = self.get_value(instr.a)
        b = self.get_value(instr.b)
        cond = eval(f"{a} {instr.cond} {b}")
        self.jump_to_block(instr.lab_yes if cond else instr.lab_no)

    def execute_jump(self, instr):
        """处理无条件跳转"""
        self.jump_to_block(instr.target)

    def execute_phi(self, instr):
        """处理PHI指令"""
        predecessor = self.current_ctx.block.previous_block
        for block, value in instr.inputs.items():
            if block == predecessor:
                self.current_ctx.variables[instr.name] = self.get_value(value)
                return
        self.current_ctx.variables[instr.name] = 0  # 默认值

    def execute_alloc(self, instr):
        # 分配内存并记录基地址
        base_address = self.next_address
        self.next_address += instr.amount
        self.allocations[instr.name] = base_address
        # 初始化内存区域为0
        for offset in range(0, instr.amount, 4):
            self.memory[base_address + offset] = 0

    def execute_address_of(self, instr):
        # 获取Alloc变量的基地址
        base_address = self.allocations.get(instr.src.name, 0)
        self.current_ctx.variables[instr.name] = base_address

    def execute_load(self, instr):
        address = self.get_value(instr.address)
        # 从内存读取值，默认0
        value = self.memory.get(address, 0)
        self.current_ctx.variables[instr.name] = value

    def execute_store(self, instr):
        value = self.get_value(instr.value)
        address = self.get_value(instr.address)
        self.memory[address] = value

    def execute_cast(self, instr):
        """处理类型转换指令"""

        def convert_value(value, from_ty, to_ty):
            """根据类型转换规则转换值"""

            # 处理指针类型转换（ptr <-> 整数）
            if isinstance(from_ty, ir.PointerTyp) and to_ty.is_integer:
                # 指针转整数：直接返回地址值（假设指针用u64表示）
                return value  # 值本身就是模拟的整数地址

            if from_ty.is_integer and isinstance(to_ty, ir.PointerTyp):
                # 整数转指针：将整数值视为地址
                return value

            # 整数类型转换逻辑
            if from_ty.is_integer and to_ty.is_integer:
                from_bits = from_ty.bits
                to_bits = to_ty.bits

                # 计算掩码保留有效位
                mask = (1 << from_bits) - 1
                truncated = value & mask

                # 有符号处理
                if from_ty.is_signed and (truncated >> (from_bits - 1)) & 1:
                    sign_extend = (0xFFFFFFFFFFFFFFFF << from_bits) | truncated
                else:
                    sign_extend = truncated

                # 根据目标类型调整符号
                if to_ty.is_signed:
                    # 符号扩展至目标位宽
                    if to_bits > from_bits:
                        return sign_extend
                    else:
                        # 截断至目标位宽并保留符号
                        return sign_extend & ((1 << to_bits) - 1)
                else:
                    # 无符号直接截断
                    return truncated & ((1 << to_bits) - 1)

            # 处理浮点数转换（暂不支持）
            if from_ty in (ir.f32, ir.f64) or to_ty in (ir.f32, ir.f64):
                raise NotImplementedError("Floating point casts not supported")

            # 默认返回原始值（如Blob转换）
            return value

        src_value = self.get_value(instr.src)
        from_ty = instr.src.ty
        to_ty = instr.ty

        # 调用类型转换函数
        converted_value = convert_value(src_value, from_ty, to_ty)

        # 存储转换结果
        self.current_ctx.variables[instr.name] = converted_value

    def execute_exit(self, instr):
        self.pop_context()

    def jump_to_block(self, new_block):
        """跳转到新块"""
        new_block.previous_block = self.current_ctx.block
        self.current_ctx.block = new_block
        self.current_ctx.ip = 0


with open('examples/quick_sort.pi', 'r') as f:
    mod = irutils.Reader().read(f)

interpreter = Interpreter()
result = interpreter.run(mod)
print(f"Execution Result: {result}")
